const fs = require('fs')
const tf = require('@tensorflow/tfjs-node')

const getData = async (trainDir , outputDir) =>{
    // 
    const classes = fs.readdirSync(trainDir).filter(n => !n.includes('.'))
    fs.writeFileSync(`${outputDir}/classes.json`, JSON.stringify(classes))


    const data = []
    classes.forEach((dir, dirIndex) =>{
        console.log(dir)

        fs.readdirSync(`${trainDir}/${dir}`)
        .filter(n => n.match(/jpg$/))
        // .slice(0,10)
        .forEach(filename => {
            // console.log('读取', dir , filename)
            const imgPath = `${trainDir}/${dir}/${filename}`
            data.push({imgPath , dirIndex})
        })
    })
    tf.util.shuffle(data)
    const ds = tf.data.generator(function* (){
        const count = data.length
        const batchSize = 32 
        for (let start = 0; start < count; start+=batchSize) {
            console.log('当前批次', start)
            const end = Math.min(start + batchSize, count);
            yield tf.tidy(() => {
                const inputs = []
                const labels = []
                for (let j = start; j < end; j++) {
                    const {imgPath, dirIndex} = data[j];
                    const x = imgPath2x(imgPath)
                    inputs.push(x);
                    labels.push(dirIndex);
                }
                
                const xs = tf.concat(inputs)
                const ys = tf.tensor(labels)
                return {
                    xs, ys
                }
            })
        }
    })

    return {
        ds, classes
    }
};
/** 将图片路径读取并转为tensor */
const imgPath2x = (imgPath) => {
    // tf.tidy 及时清理转换后的内存
    const buffer =  fs.readFileSync(imgPath)
    return tf.tidy(() => {
        const imgTs = tf.node.decodeImage(new Uint8Array(buffer))
        const imgTsResized = tf.image.resizeBilinear(imgTs, [224,224])
        // 将归一化后的结果0 -255 压缩到-1 - 1 之间， 并且转换为模型需要的形状 224*224 3代表彩色 1代表 拓展一维
        return imgTsResized.toFloat().sub(255/2).div(255/2).reshape([1,224,224,3])
    })
   
}

module.exports = getData